import torch
import matplotlib.pyplot as plt
import random


'''
    1.数据处理
    2.构造模型
    3.数据训练
    4.图形可视化
'''


def create_data(w,b,data_num):
    '''
    生成简单的线性回归模型的数据集
    :param w:真实的权重向量，形状为（4,）的张量
    :param b:真实的偏置值,形状为标准张量
    :param data_num:数据点的数量，取值为500（int）
    :return:返回生成的输入数据x和输出数据y
    '''
    x = torch.normal(0,1,(data_num,len(w))) #生成一个形状为（500，4）的张量x，其中元素从均值为0，标准差为1的正态分布中得到
    y = torch.matmul(x,w) + b #matmul表示矩阵相乘，生成一个形状为（500，1）的张量y 表达式为 y = x * w + b


    #生成一个形状为（500，1）的张量noise，其中元素从均值为0，标准差为0.01的正态分布中得到
    noise = torch.normal(0,0.01,y.shape) # xxx.shape表示生成与xxx相同的张量形状，
    y += noise #y加上噪声对其的影响
    return x,y #返回值


num = 500 #训练集的大小；500


true_w = torch.tensor([8.1,2,2,4]) #真实得到的w
true_b = torch.tensor(1.1) #真实的到的b

X,Y = create_data(true_w,true_b,num) #调用函数，得到生成的X,Y
#用于绘制散点图
#X[:,3]:表示取x张量中的第四列
#Y:由于Y是（500，1）的张量，故只需要取这一列即可
#1：表示绘制的点的大小，这里大小为1
plt.scatter(X[:,3],Y,1) #
plt.show() #用于显示绘制的图形
def data_provider(data,label,batchsize):
    '''
    每次访问这个函数，就能提供一批数据
    :param data:输入数据，是一个二维张量，形状为 (data_num, features)，表示有 data_num（500） 个样本，每个样本有 features（4） 个特征
    :param label:目标标签，是一个一维张量，形状为 (data_num（500）,)，表示每个样本的目标值
    :param batchsize:批次，每个批次的数据量，表示每次返回的数据点数量
    :return:无返回值
    '''
    length = len(label) #计算目标标签的长度，即数据集的点数
    indices = list(range(length)) #生成一个0~500(左闭右开)的列表
    random.shuffle(indices) #使用 random.shuffle 随机打乱索引列表，确保数据在每个批次中是随机的
    for each in range(0,length,batchsize): #遍历所有数据点，每次取 batchsize 个数据点。
        get_indices = indices[each:each+batchsize] #get_indices 是当前批次的索引列表。
        get_data = data[get_indices] #get_data 是当前批次的输入数据，通过 data[get_indices] 获取
        get_label = label[get_indices] #是当前批次的目标标签，通过 label[get_indices] 获取

        yield get_data,get_label #生成器，返回当前批次的数据和标签，每次调用次函数返回同一批的生成数据和生成标签

batchsize = 16 #批次设置为16


def fun(x,w,b):
    '''
    线性回归模型的预测函数
    :param x:输入数据，是一个二维张量，形状为 (data_num, features)，表示有 data_num 个样本，每个样本有 features 个特征
    :param w:权重向量，是一个一维张量，形状为 (features,)，表示每个特征的权重
    :param b:偏置项，通常是一个标量张量，表示线性模型的偏置项
    :return:返回计算得到的预测值 pred_y
    '''
    pred_y = torch.matmul(x,w) + b # 使用 torch.matmul 计算输入数据 x 和权重向量 w 的矩阵乘法，然后加上偏置项 b，得到预测值 pred_y
    return pred_y #返回偏置值


def maeLoss(pred_y,y):
    '''
    计算平均误差
    :param pred_y: 模型预测的目标值，是一个一维张量，形状为 (data_num,) 或 (data_num, 1)，表示有 data_num 个预测值
    :param y:真实的目标值，是一个一维张量，形状为 (data_num,) 或 (data_num, 1)，表示有 data_num 个真实值
    :return:
    '''
    return torch.sum(abs(pred_y - y)) / len(y) #返回计算得到的平均绝对误差


def sgd(paras,lr): #Stochastic Gradient Descent（随机梯度下降）
    '''
    随机梯度下降，更新参数
    :param paras: 模型的参数列表，通常包括权重和偏置项
    :param lr: 学习率，控制参数更新的步长
    :return:
    '''
    with torch.no_grad(): #确保在更新参数时不会计算梯度。这是因为参数更新是一个非训练步骤，不需要计算梯度
        for para in paras:
            para -= para.grad * lr #不能写成para = para - para.grad * lr，因为会创建一个新的张量，导致原有的张量更新失败
            para.grad.zero_() #使用过梯度归0

lr = 0.03 #学习率，控制参数更新的步长
w_0 = torch.normal(0,0.01,true_w.shape,requires_grad=True)
#初始化权重向量，从均值为0、标准差为0.01的正态分布中随机抽取，形状与 true_w 相同，并设置 requires_grad=True 以便计算梯度
b_0 = torch.tensor(0.01,requires_grad=True)
#初始化偏置项，设置为0.01，并设置 requires_grad=True 以便计算梯度
# print(w_0,b_0)
epochs = 50 #训练的总轮数，表示整个数据集会被遍历50次


for epoch in range(epochs):
    data_loss = 0 #初始化为0，用于记录当前轮次的总损失
    for batch_x, batch_y in data_provider(X, Y, batchsize):
        pred_y = fun(batch_x,w_0,b_0) #调用预测函数 fun，计算当前批次的预测值
        loss =maeLoss(pred_y,batch_y) #调用 MAE 损失函数 maeLoss，计算当前批次的损失。
        loss.backward() #调用 backward 方法，梯度回传
        sgd([w_0,b_0],lr) #调用 SGD 优化器函数 sgd，更新参数 w_0 和 b_0
        data_loss += loss #累积损失
    print('epoch %03d: loss: %.6f' % (epoch,data_loss))

print('真实的函数值是',true_w,true_b)
print('深度训练50次得到的参数值是',w_0,b_0)

idx = 2 # 选择第3个特征（索引为2）进行可视化。
plt.plot(X[:,idx].detach().numpy(),X[:,idx].detach().numpy()*w_0[idx].detach().numpy()+b_0.detach().numpy())
plt.scatter(X[:,idx],Y,1)
plt.show()